#pragma once

#include "limbo/ip/ipv4_endpoint.h"
#include "state.hpp"

namespace limbo::udp {

template <typename UdpState> struct Receiver {
  using ContextState = Receiver;
  using Context = Receiver;
  using LowerState = UdpState;
  using LowerContext = typename UdpState::Context;
  using LowerPacket = typename LowerState::Packet;
  using IpState = typename UdpState::LowerState;
  using IpContext = typename IpState::Context;

  inline void init(const ip::IPv4Endpoint &source, uint8_t ttl = 50) {
    ip_ctx.source = source.address;
    ip_ctx.proto = ip::Proto::udp;
    ip_ctx.ttl = ttl;
    udp_ctx.source_port = source.port;
  }

  Result recv(Chunk, LowerState *lower_state) {
    state = lower_state;
    return Result::make_consumed(this, state->get_parsed().payload);
  }

  LowerPacket &get_parsed() { return state->get_parsed(); }

  LowerContext *get_context() { return &udp_ctx; }

  inline bool select(const LowerPacket *packet) {
    return (packet->dest_port == udp_ctx.source_port);
  }

  template <typename Stack>
  Result send(Stack &stack, Chunk buff, const ip::IPv4Endpoint &dest,
              Chunk payload) {
    ip_ctx.destination = dest.address;
    udp_ctx.dest_port = dest.port;
    return stack.send(buff, *this, payload);
  }

private:
  IpContext ip_ctx;
  LowerContext udp_ctx = {0, 0, &ip_ctx};
  LowerState *state;
};

template <typename State> inline auto get_lower_context(Receiver<State> &ctx) {
  return ctx.get_context();
}

template <typename State> inline auto &get_ip_context(Receiver<State> &ctx) {
  auto udp_ctx = ctx.get_context();
  auto ip_ctx = get_lower_context(*udp_ctx);
  return *ip_ctx;
}

template <typename State> inline auto &get_link_context(Receiver<State> &ctx) {
  using IpState = typename Receiver<State>::IpState;
  static_assert(
      details::has_lower_state_v<IpState>,
      "the function has meaning only when there is link state in the stack");
  auto udp_ctx = ctx.get_context();
  auto ip_ctx = get_lower_context(*udp_ctx);
  auto link_ctx = get_lower_context(*ip_ctx);
  return *link_ctx;
}

} // namespace limbo::udp
